"""
Library module for code shared between rankAbundance, rarefactionCurve and
others. Contains code for selecting random subsamples from full datasets.
"""

from ssummo.traverse import dict_walk
from ssummo.ssummolib import collapse_at_rank

from math import log, fsum
import random
import sys
from threading import Thread

class Indices:
    """Number crunching class. Methods calculate biodiversity statistics for
    SSuMMo results files."""
    #def __init__(self, options=None):
    def __init__(self, out, collapse_rank=None):
        """
        """
        #if options is None:
        #    self.options = LocalOptions([])
        #else:
        #    self.options = options
        self.out = out          #: Output file
        self.shannon = None     #: Shannon value
        self.jack = None        #: Jackknife value
        self.simpson = None     #: Simpson index
        self.name = None        #: Dataset name
        self.collapse_rank = collapse_at_rank
        #: To trim tree at specified rank before resampling.

    def calc_all(self):
        """Automates all calculations so that we can print the output.
        """
        counts, probs, total = self.tally( )
        self.richness, max_evenness = self.shannon_index(probs)
        #diversity = self.simpsons_index(counts, total)
        self.jack = self.jackknife(counts, total)

    def jackknife(self, counts, total):
        k = counts.count(1)  # number of 'unique' species.
        n = len(counts)      # number of 'plots' - observed taxa
        S = n + ((n - 1.) / n) ** k  # jackknife estimate.
        self.jack = S
        return S

    def load(self, results_dictionary, combined=False, name=None):
        """Load the dictionary. If the option '-collapse-at-rank' is given,
        then we shall also collapse the dictionary at the specificied rank.
        Results dictionary is not returned, but is stored as:-
        self.dictionary
        """
        if self.collapse_rank is not None:
            results_dictionary = collapse_at_rank(results_dictionary,
                self.collapse_rank, 
                #self.options['-collapse-at-rank'],
                combined=combined,
                in_files=name,
                #in_files=self.options['-in']
                )
        #if self.options['-collapse-at-rank'] is not None:
            #results_dictionary = collapse_at_rank(results_dictionary, self.options, combined=combined)
        self.dictionary = results_dictionary
        if name is not None:
            self.name = name

    def print_header( self ):
        print('#{0} {1} {2} {3} {4}'.format(
              ' File'.ljust(25),
              'Shannon'.ljust(20) ,
              'Max Evenness'.ljust(25),
              'Simpsons'.ljust(25),
              'Jackknife'.ljust(25), ))
        print('#{0} {1} {2} {3} {4}'.format(
              ' name'.ljust(25),
              '(H`)'.ljust(20) ,
              '(Hmax)'.ljust(25),
              'Richness'.ljust(25),
              '(sampling bias)'.ljust(25), ))

    def __str__(self):
        """If you haven't already calculated everything and haven't specified
        :data:`.self.out` as anything, then we'll have to recalculate
        everything and reprint the headers for good measure..."""
        if (self.shannon and self.jack and self.simpson) is None:
            self.calc_all()
        #    if self.options['-out'] == sys.stdout:
        #        self.__init__( self.options )
        if self.name is None:
            self.name = '<unknown>'
        return '{0}  {1} {2} {3} {4}'.format(
                    self.name,
                    '{0:.5f}'.format(self.shannon[0]).ljust(20),
                    '{0:.5f}'.format(self.shannon[1]).ljust(25),
                    '{0:.5f}'.format(self.simpson   ).ljust(25),
                    '{0:.5f}'.format(self.jack      ).ljust(25))

    def tally(self, combined=False):
        """Writes rank abundance information to :data:`.self.out`.
        This is represented as just a single column, showing the number
        of sequences allocated to each taxon.

        To use this, first define self.dictionary, preferably with
        :meth:`Indices.load`

        The function returns 3 items:-

        1) counts-a sorted list, containing the number of sequences
           allocated to each taxa, probabilities and total number of
           taxa.
        2) probs - a list of floating point numbers, showing the percentage
           of sequences assigned to each taxon.
        3) total - the total number of sequences found in the dataset.

        These 3 items are needed to calculate the other indices...
        """
        counts = []
        for path, node in dict_walk('', self.dictionary ):
            if 'accessions' in node:
                if combined:
                    count = 0
                    for dataset in node['accessions']:
                        for prob in dataset:
                            if float( prob ) > 0.:
                                count += 1
                    counts.append( count )
                elif len( node['accessions'] ) == 0:
                    continue
                else:
                    counts.append(len(node['accessions']))
        if self.out != sys.stdout and not isinstance(self.out, file):
            self.out = file(self.out, 'w')
        counts = sorted(counts, reverse=True)
        total = sum(counts )

        probs = [ float(count) / total for count in counts ]
        self.out.write('\n'.join([str(100. * prob) for prob in probs]))
        self.out.flush()
        #self.options['-out'].write( '\n'.join( [ str(100.*prob) for prob in probs ] ))
        #self.options['-out'].flush()
        return counts, probs, total

    def shannon_index(self, probs):
        """Give the probabilities as returned from tally() and total
        species count.
        :returns: The tuple :math:`(H', H_{max}, )`
        """
        H_prime = - fsum( [ p * log(p) for p in probs ] )
        H_max = log( len(probs) )
        self.shannon = (H_prime, H_max, )
        return self.shannon

    def simpsons_index(self, counts, total):
        """Simpson's index of Diversity.
        Returns a value from 0 - 1, where 0 shows no
        diversity, and 1 shows infinite diversity.
        """
        diversity = 1. - (fsum([n * (n - 1.) for n in counts]) /
                          (total * (total - 1.)))
        self.simpson = diversity
        return diversity

    def SPlusJackknife(self, counts, total):
        """Implementation of JackKnife formula from:
        http://www.jstatsoft.org/v15/i03/paper

        Smith, CD. Pontius, JS. Jackknife Estimator of Species Richness with
        S-PLUS, 2006, 15 (3)

        .. math::

            J_{n}(S) = S_{0} + \\frac{(n - 1)}{n} \cdot \sum_{i=1}^{n}{r_{i}}
        """
        #S0 = len(counts)
        n = len(counts)
        self.jack = JnS = total + ((n - 1.) / n) * sum(counts)
        #print('S0: ', S0)
        print('n: ', n)
        print('Jns: ', JnS)
        return JnS

class RandomNumberThread(Thread):
    def __init__(self, Q, N, replicates, increment):
        """Need a Queue.Queue object, number of accessions
        which we're rarefying, the number of replicates for
        each point, and the increment between points.
        """
        Thread.__init__(self)
        self.N = N
        self.Q = Q
        self.increment = increment
        self.replicates = replicates

    def run( self ):
        i = self.increment
        while i < self.N:
            for replicate in xrange( self.replicates ):
                n_left = self.N
                for ii in xrange( i ):
                    rand = random.randint( 0, n_left-1 )
                    self.Q.put_nowait( rand )
                    n_left -= 1
            i += self.increment
        self.Q.put( 'done' )
        self.Q.join()
        return

class Replicate():
    """Store statistics on a single dataset. Supports addition and division."""

    __slots__ = ['nseqs', 'rarefaction', 'shannon', 'shannon_max', 'simpson']

    nseqs = 0
    #: number of sequences sampled
    rarefaction = float()
    #: 
    shannon = float()
    #: Shannon :math:`H'` value
    shannon_max = float()
    #: Shannon :math:`H_{max}` value
    simpson = float()
    #: Simpson biodiversity index.

    def __init__(self):
        """Create new, empty replicate."""
        #self.nseqs = int()
        #for slot in self.__slots__:
        #    setattr(self, slot, float())

    def __add__(self, other):
        new = Replicate()
        for slot in self.__slots__:
            setattr(new, slot, getattr(self, slot) + getattr(other, slot))
        return new

    def __div__(self, other):
        new = Replicate()
        for slot in self.__slots__:
            setattr(new, slot, getattr(self, slot) / getattr(other, slot))
        return new

class Results:
    """Object for storing multiple :class:`Replicate` instance statistics."""

    __slots__ = ['file_name', 'xvals', 'rarefaction', 'simpson', 'shannon',
                 'shannon_max', 'jack']

    def __init__(self):
        self.file_name = str()
        #: dataset name
        self.xvals = list()
        #: number of random samples (sequences)
        self.rarefaction = dict()
        #: number of observed OTUs (OTU defined by rank)
        self.simpson = dict()
        #: simpson index, as returned by `.Indices.simpsons_index`
        self.shannon = dict()
        #: shannon index, as returned by `.Indices.shannon_index`
        self.shannon_max = dict()
        #: shannon_max as returned by `.Indices.shannon_index`
        self.jack = dict()
        #: Once all the replicates have been calculated, `Results.jackknife`
        #: can be called, which will jackknife the replicates.

    def __get__(self, key):
        return getattr(self, key)

    def __getitem__(self, key):
        rep = Replicate()
        rep.nseqs = self.xvals[key]
        rep.rarefaction = self.rarefaction[key]
        rep.simpson = self.simpson[key]
        rep.shannon = self.shannon[key]
        rep.shannon_max = self.shannon_max[key]
        rep.jack = self.jack[key]
        return rep

    def __iter__(self):
        for i in xrange( len( self.xvals ) ):
            yield self[i]

    def append(self, replicate):
        """Append new replicate data to this :class:`.Results` instance."""
        if not isinstance(replicate, Replicate):
        #if type( replicate ) != Replicate:
            raise AttributeError( "replicate must be a Replicate instance" )
        xval = replicate.nseqs
        if xval not in self.xvals:
            self.xvals.append( xval )
            self.rarefaction.update( { xval : [] } )
            self.simpson.update( { xval : [] } )
            self.shannon.update( { xval : [] } )
            self.shannon_max.update( { xval : [] } )
            self.jack.update( { xval : [] } )
        self.rarefaction[xval].append( replicate.rarefaction )
        self.simpson[xval].append( replicate.simpson )
        self.shannon[xval].append( replicate.shannon )
        self.shannon_max[xval].append( replicate.shannon_max )
        self.jack[xval].append( replicate.jack )

    def jackknife(self):
        """Calculate jackknife information using stored result data."""
        n = len(self.nseqs)
        i = 0
        theta = sum(self.jack)
        thetas = []
        while i < n:
            theta_i = theta - self.jack[i]
            theta_hat = n * theta - (n - 1) * theta_i
            thetas.append(theta_hat)
            i += 1
        first_order = 1 / n * sum(thetas)
        self.jack = first_order
        return first_order
